{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6ba5137f",
   "metadata": {},
   "source": [
    "# Mean and Variance Estimates for Multiple Randomization Designs\n",
    "\n",
    "In a two-sided marketplace, multiple sellers and multiple buyers interact. For instance, on an e-commerce platform,  which is a typical type of a two-sided marketplace, buyers see a sellers' offers and choose to buy or not. The platform wants to assess a new info policy's impact on a buyer's likelihood of buying an item. In a traditional A/B test, researchers can randomize buyers to A (treatment group) or B (control group) or sellers to A or B. These single randomization designs are traditional RCTs, and can lead to biased average treatment effect estimate. \n",
    "\n",
    "For example, the platform wants to recommend more TVs during black Friday to customers, and wants to test if the strategy is effective using A/B test. If researchers carry out an user side experiments, users in treatment and control group will compete for the same product. If treatment group users buy all TVs, which leads to control group users have nothing to buy, average treatment effect can be over-estimate. Therefore a two-sided A/B test is desired. \n",
    "\n",
    "We follow Bajari, Patrick, et al. (2021) (https://arxiv.org/pdf/2112.13495) formulas and provide treatment effect estimators for the following types.  \n",
    "\n",
    "\n",
    "### Estimators\n",
    "\n",
    "We have following types of estimators:\n",
    "* $\\tau(p^B,p^S)$: the treatment effect of treated buyer and treated seller pairs vs control buyer and control seller pair. The estimator is $\\bar{\\bar{Y_t}}-\\bar{\\bar{Y_{c}}}$, variance is \n",
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "V(\\tau(p^B, p^S)) = V(\\bar{\\bar{Y_t}}-\\bar{\\bar{Y_{c}}})=\n",
    "&V(\\bar{\\bar{Y_t}})+V(\\bar{\\bar{Y_{c}}})-2C(\\bar{\\bar{Y_t}},\\bar{\\bar{Y_{c}}})\\\\\n",
    "\\end{aligned}\n",
    "\\end{equation}\n",
    "* $\\tau_{direct}$: the treatment effect of treated buyer and treated seller pairs vs control buyer and control seller pair, minus spill-over effect. The estimator is $\\bar{\\bar{Y_t}}-\\bar{\\bar{Y_{ib}}}-\\bar{\\bar{Y_{is}}}+\\bar{\\bar{Y_{c}}}$, variance is \n",
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "V(\\tau_{direct}) = V(\\bar{\\bar{Y_t}}-\\bar{\\bar{Y_{ib}}}-\\bar{\\bar{Y_{is}}}+\\bar{\\bar{Y_{c}}})=\n",
    "&V(\\bar{\\bar{Y_t}})+V(\\bar{\\bar{Y_{ib}}})+V(\\bar{\\bar{Y_{is}}})+V(\\bar{\\bar{Y_{c}}})\\\\\n",
    "&-2C(\\bar{\\bar{Y_t}},\\bar{\\bar{Y_{ib}}})-2C(\\bar{\\bar{Y_t}},\\bar{\\bar{Y_{is}}})+2C(\\bar{\\bar{Y_t}},\\bar{\\bar{Y_{c}}})\\\\\n",
    "&+2C(\\bar{\\bar{Y_{ib}}},\\bar{\\bar{Y_{is}}})-2C(\\bar{\\bar{Y_{ib}}},\\bar{\\bar{Y_{c}}})\\\\\n",
    "&-2C(\\bar{\\bar{Y_{is}}},\\bar{\\bar{Y_{c}}}))\\\\\n",
    "\\end{aligned}\n",
    "\\end{equation}\n",
    "* $\\tau^b_{spillover}$: spill-over effect from treated users to control users. The estimator is $\\bar{\\bar{Y_{ib}}}-\\bar{\\bar{Y_{c}}}$, variance is \n",
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "V(\\tau^b_{spillover}) = V(\\bar{\\bar{Y_{ib}}}-\\bar{\\bar{Y_{c}}})=\n",
    "&V(\\bar{\\bar{Y_{ib}}})+V(\\bar{\\bar{Y_{c}}})-2C(\\bar{\\bar{Y_{ib}}},\\bar{\\bar{Y_{c}}})\\\\\n",
    "\\end{aligned}\n",
    "\\end{equation}\n",
    "* $\\tau^s_{spillover}$: spill-over effect from treated sellers to control sellers. The estimator is $\\bar{\\bar{Y_{is}}}-\\bar{\\bar{Y_{c}}}$, variance is \n",
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "V(\\tau^s_{spillover}) = V(\\bar{\\bar{Y_{is}}}-\\bar{\\bar{Y_{c}}})=\n",
    "&V(\\bar{\\bar{Y_{is}}})+V(\\bar{\\bar{Y_{c}}})-2C(\\bar{\\bar{Y_{is}}},\\bar{\\bar{Y_{c}}})\\\\\n",
    "\\end{aligned}\n",
    "\\end{equation}\n",
    "\n",
    "In formulas above, $V(\\bar{\\bar{Y_t}})$, $V(\\bar{\\bar{Y_{ib}}})$, $V(\\bar{\\bar{Y_{is}}})$, and $V(\\bar{\\bar{Y_c}})$ follows Lemma A.4. Covariance terms follows Lemma A.5.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9708d3f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['/Library/Frameworks/Python.framework/Versions/3.12/lib/python312.zip', '/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12', '/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/lib-dynload', '', '/Users/bytedance/Library/Python/3.12/lib/python/site-packages', '/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages']\n",
      "current version is:  0.0.5\n"
     ]
    }
   ],
   "source": [
    "# import sys\n",
    "# print(sys.path)\n",
    "# sys.path.append('/Users/bytedance/PycharmProjects/github/CausalMatch')\n",
    "\n",
    "import causalmatch as causalmatch\n",
    "from causalmatch import matching, gen_test_data, gen_test_data_mrd\n",
    "from causalmatch import mrd\n",
    "\n",
    "print('current version is: ',causalmatch.__version__)\n",
    "\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import statsmodels.api as sm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4a6ca681",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>shop_id</th>\n",
       "      <th>user_id</th>\n",
       "      <th>treatment</th>\n",
       "      <th>treatment_u</th>\n",
       "      <th>treatment_s</th>\n",
       "      <th>error</th>\n",
       "      <th>status</th>\n",
       "      <th>y_clean</th>\n",
       "      <th>y_overflow</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.046911</td>\n",
       "      <td>t</td>\n",
       "      <td>1.046911</td>\n",
       "      <td>2.046911</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.028286</td>\n",
       "      <td>ib</td>\n",
       "      <td>-0.028286</td>\n",
       "      <td>0.471714</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>-0.150906</td>\n",
       "      <td>t</td>\n",
       "      <td>0.849094</td>\n",
       "      <td>1.849094</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>-0.113563</td>\n",
       "      <td>t</td>\n",
       "      <td>0.886437</td>\n",
       "      <td>1.886437</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.121211</td>\n",
       "      <td>ib</td>\n",
       "      <td>0.121211</td>\n",
       "      <td>0.621211</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   shop_id  user_id  treatment  treatment_u  treatment_s     error status  \\\n",
       "0        1        1          1            1            1  0.046911      t   \n",
       "1        2        1          0            1            0 -0.028286     ib   \n",
       "2        3        1          1            1            1 -0.150906      t   \n",
       "3        4        1          1            1            1 -0.113563      t   \n",
       "4        5        1          0            1            0  0.121211     ib   \n",
       "\n",
       "    y_clean  y_overflow  \n",
       "0  1.046911    2.046911  \n",
       "1 -0.028286    0.471714  \n",
       "2  0.849094    1.849094  \n",
       "3  0.886437    1.886437  \n",
       "4  0.121211    0.621211  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# STEP 1: generate synthetic data of a two-sided experiment:\n",
    "#         In the test dataset, shop_id mimics seller id in Bajari, Patrick, et al. (2021),\n",
    "#         user_id mimics buyer id, treatment_u is the treatment status of the user_id,\n",
    "#         treatment_s is the treatment status of the shop_id, y_overflow is the dependent variable\n",
    "#         target to be estimated. \n",
    "\n",
    "df_raw = gen_test_data_mrd(n_shops = 1000\n",
    "                           , n_users = 10\n",
    "                           , ate   = 1 # ATE\n",
    "                           , uflow = 0.5 # user side overflow effect\n",
    "                           , sflow = 0.5 # seller side overflow effect\n",
    "                           , fraction_bt = 0.5 # treatment group fraction\n",
    "                           , fraction_st = 0.5 # control group fraction\n",
    "                           , seed = 123456) # random seed number\n",
    "df_raw.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a7027200",
   "metadata": {},
   "outputs": [],
   "source": [
    "# STEP 2: initialize a mrd object\n",
    "mrd_obj = mrd(data = df_raw,\n",
    "              idb = 'user_id',\n",
    "              ids = 'shop_id',\n",
    "              tb  = 'treatment_u',\n",
    "              ts  = 'treatment_s',\n",
    "              y   = 'y_overflow')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "85510f58",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>parameters</th>\n",
       "      <th>mean</th>\n",
       "      <th>variance</th>\n",
       "      <th>t_stat</th>\n",
       "      <th>p_values</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>tau</td>\n",
       "      <td>2.000219</td>\n",
       "      <td>1.279201</td>\n",
       "      <td>1.768512</td>\n",
       "      <td>0.077052</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>tau_tdirect</td>\n",
       "      <td>0.999174</td>\n",
       "      <td>0.284924</td>\n",
       "      <td>1.871873</td>\n",
       "      <td>0.061297</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>tau_b_spillover</td>\n",
       "      <td>0.501487</td>\n",
       "      <td>0.079940</td>\n",
       "      <td>1.773693</td>\n",
       "      <td>0.076190</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>tau_s_spillover</td>\n",
       "      <td>0.499558</td>\n",
       "      <td>0.023354</td>\n",
       "      <td>3.268926</td>\n",
       "      <td>0.001089</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        parameters      mean  variance    t_stat  p_values\n",
       "0              tau  2.000219  1.279201  1.768512  0.077052\n",
       "1      tau_tdirect  0.999174  0.284924  1.871873  0.061297\n",
       "2  tau_b_spillover  0.501487  0.079940  1.773693  0.076190\n",
       "3  tau_s_spillover  0.499558  0.023354  3.268926  0.001089"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# STEP 3: calculate average treatment effect of four type\n",
    "mrd_obj.ate2()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
